import os
import cv2
import subprocess
import random



def run_fog_test(test_strength, model_name):
    model = f'../{model_name}.pt'
    
    update_testsets = f" \
    rm ../../datasets/kitti/images/test/* &&\
    cp ../../datasets/fogged/fogged_strength{test_strength}/* ../../datasets/kitti/images/test/ && \
    echo 'Test set updated successfully!' \
    " 
    subprocess.run(update_testsets, shell=True, check=True)
    
    val_command = f" \
    python ../val.py \
    --data ../data/kitti.yaml \
    --weights {model} \
    --task test &&\
    echo 'Test set val successfully!' \
    " 
    subprocess.run(val_command, shell=True, check=True)
    

init_ratio = {
    '0.6':0,
}
# 这个是在当前文件夹的父目录调用的，所以路径是这样的。
def mix_dataset(fogged_folder = '../datasets/fogged/', 
                ratio = init_ratio,
                train_folder = '../datasets/kitti/images/origin_train', 
                out_folder = '../datasets/kitti/images/train'
               ):
    
    if sum([value for key,value in init_ratio.items()]) > 1.0:
        print("ratio error")
        return 
    all_subdirs = [d for d in os.listdir(fogged_folder) if os.path.isdir(os.path.join(fogged_folder, d))]
    test_strengths = [d.split('train_fogged_strength')[1] for d in all_subdirs if 'train_fogged_strength' in d]
    
    img_names = [img for img in os.listdir(train_folder) if img.endswith('.png')]
    img_num = len(img_names)

    random.shuffle(img_names)
    begin = 0
    end = 0;
    for key, value in ratio.items():
        end = int(begin + value * img_num)
        for img_name in img_names[begin : end]:
            cp_command = f" \
                cp {fogged_folder}train_fogged_strength{key}/{img_name} {out_folder} \
            " 
            subprocess.run(cp_command, shell=True, check=True)
        begin = end
    for img_name in img_names[begin : ]:
        cp_command = f" \
            cp {train_folder}/{img_name} {out_folder} \
        " 
        subprocess.run(cp_command, shell=True, check=True)
    
        
    
    

